#!/bin/bash

set -e

TOP_DIR=$(cd "$(dirname "$0")"/.. || exit; pwd)
OHOS_IMG="${TOP_DIR}/images"
FW_PAYLOAD_BIN="${OHOS_IMG}/fw_payload.bin"

function usage()
{
    echo "usage:  <disk>"
    echo "  disk: target disk to burn"
    echo "  example: $0 /dev/sdb"
}

# Function to check if a file exists
function file_exists() {
    if [ ! -f "$1" ]; then
        echo "File $1 does not exist."
        echo "Please put this burn tool and the images in the \"images\" directory."
        exit 1
    fi
}

# Function to get file size in bytes
function get_file_size() {
    stat -c %s "$1"
}

# Function to convert bytes to MiB. round to MiB
function bytes_to_mib() {
    awk "BEGIN {print int(($1 + 1024 * 1024) / (1024 * 1024))}"
}

# Function to create partitions based on image sizes
function create_partitions() {
    local start=0
    local end=0
    local current_partition=3
    local total_size=0
    local first_partition_size=$((15 * 1024 * 1024))
    local second_partition_size=$((112 * 1024 * 1024))
    local target_disk="$1"
    
    # Calculate the total size needed for all images
    for img in ${OHOS_IMG}/system.img ${OHOS_IMG}/vendor.img ${OHOS_IMG}/userdata.img ${OHOS_IMG}/chip_prod.img ${OHOS_IMG}/eng_system.img ${OHOS_IMG}/sys_prod.img ${OHOS_IMG}/updater.img; do
        total_size=$((total_size + $(get_file_size "$img")))
    done
    
    # Add extra space for the first partition
    total_size=$((total_size + first_partition_size))
    
    # Convert total size to MiB and add some buffer
    total_size_mib=$(bytes_to_mib "$total_size")
    total_size_mib=$((total_size_mib + 100))

    # Clear existing partition table and create a new one
    sgdisk --clear -g ${target_disk}
    
    # Create the first partition for FW_PAYLOAD_BIN
    echo "Creating partition for bootloader..."
    end=$((start + $(bytes_to_mib $first_partition_size)))
    echo sgdisk --new 1:$start:${end}Mib -u=1:867064e6-49ad-33bd-8346-1a1eb81ade18 ${target_disk}
    sgdisk --new 1:$start:${end}Mib -u=1:867064e6-49ad-33bd-8346-1a1eb81ade18 ${target_disk}

    echo "Burning ${FW_PAYLOAD_BIN} to ${target_disk}1..."
    dd if=${FW_PAYLOAD_BIN} of=${target_disk}1 status=progress oflag=sync bs=1M

    start=$((end + 1))

    # Create the second partition for boot
    echo "Creating partition for boot..."
    end=$((start + $(bytes_to_mib $second_partition_size)))
    echo sgdisk --new 2:${start}Mib:${end}Mib ${target_disk}
    sgdisk --new 2:${start}Mib:${end}Mib ${target_disk}
    start=$((end + 1))

    # Create partitions for each image
    for img in ${OHOS_IMG}/system.img ${OHOS_IMG}/vendor.img ${OHOS_IMG}/userdata.img ${OHOS_IMG}/chip_prod.img ${OHOS_IMG}/eng_system.img ${OHOS_IMG}/sys_prod.img ${OHOS_IMG}/updater.img; do
        echo "Creating partition for $img..."
        end=$((start + $(bytes_to_mib $(get_file_size "$img"))))
        echo sgdisk --new $current_partition:${start}Mib:${end}Mib ${target_disk}
        sgdisk --new $current_partition:${start}Mib:${end}Mib ${target_disk}

        echo "Burning $img to ${target_disk}$current_partition..."
        dd if=${img} of=${target_disk}${current_partition} status=progress oflag=sync bs=10M

        start=$((end + 1))
        ((current_partition++))
    done
    
    # Uncomment below if create the last partition using the remaining space
    # sgdisk --new $current_partition:$start:$total_size_mib ${target_disk}
}

function post_burn()
{
    target_disk="$1"
    tmpdir=$(mktemp -d)
    trap 'rm -rf "$tmpdir"' EXIT
    mkfs.ext4 ${target_disk}2
    mount ${target_disk}2 ${tmpdir}
    mkdir -p ${tmpdir}/boot
    cp ${OHOS_IMG}/Image.gz ${tmpdir}/boot/
    cp ${OHOS_IMG}/a1.dtb ${tmpdir}/boot/ultrarisc-a1.dtb
    umount ${tmpdir}
}

# burn the image to the target disk
function make_burn()
{
    target_disk="$1"
    
    # Check if the target disk is specified
    if [ -z "$target_disk" ]; then
        usage
        exit 1
    fi
    
    # Verify that the target disk exists
    if [ ! -b "$target_disk" ]; then
        echo "Target disk $target_disk does not exist."
        exit 1
    fi

    # Check if all required files exist
    for img in ${FW_PAYLOAD_BIN} ${OHOS_IMG}/system.img ${OHOS_IMG}/vendor.img ${OHOS_IMG}/userdata.img ${OHOS_IMG}/chip_prod.img ${OHOS_IMG}/eng_system.img ${OHOS_IMG}/sys_prod.img ${OHOS_IMG}/updater.img; do
        file_exists "$img"
    done

    # Create partitions
    echo "Creating partitions..."
    create_partitions ${target_disk}

    echo "Post burn operations..."
    post_burn ${target_disk}

    echo "Burning completed."
}

make_burn "$1"
